# -*- coding: utf-8 -*-
import time

from langchain_community.embeddings import SparkLLMTextEmbeddings
from milvus_model import DefaultEmbeddingFunction
from pymilvus import connections, db, model
from pymilvus import MilvusClient, DataType
import pandas as pd
import os
import numpy as np
from embedding import demo
from logs import logger as log

# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
# embedding_fn:DefaultEmbeddingFunction = model.DefaultEmbeddingFunction()

_URI = "http://117.72.42.129:19530"
_TOKEN = "root:ycrqwe"
_DB_NAME = "tender"


# connections.connect(host="101.43.59.219", port=19530)


# database = db.create_database("book")
# print(database)

def t1():
    db.using_database("book")
    list = db.list_database()
    print(list)


def get_connect():
    connections.connect(
        uri=_URI,
        token=_TOKEN,
        db_name=_DB_NAME
    )


def get_client():
    client = MilvusClient(
        uri=_URI,
        user=_TOKEN.split(':')[0],
        password=_TOKEN.split(':')[1],
        db_name=_DB_NAME
    )
    return client


def create_connection():
    # 1. Set up a Milvus client
    client = get_client()
    # 2. Create a collection in quick setup mode
    client.create_collection(
        collection_name="notice",
        dimension=768,
    )
    res = client.get_load_state(
        collection_name="notice"
    )
    print(res)


def custom_schema():
    # 3. Create a collection in customized setup mode
    # 3.1. Create schema
    client = get_client()
    schema = client.create_schema(
        auto_id=True,
        enable_dynamic_field=True,
    )
    # 3.2. Add fields to schema
    schema.add_field(field_name="notice_id", datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="project_name", datatype=DataType.VARCHAR, max_length=128)
    schema.add_field(field_name="project_no", datatype=DataType.VARCHAR, max_length=64)
    schema.add_field(field_name="question", datatype=DataType.VARCHAR, max_length=128)
    schema.add_field(field_name="question_vector", datatype=DataType.FLOAT_VECTOR, dim=768)
    schema.add_field(field_name="answer", datatype=DataType.VARCHAR, max_length=65535)
    # 3.3 prepare index parameters
    index_params = client.prepare_index_params()
    # 3.4 add index
    index_params.add_index(
        field_name="notice_id",
        index_type="STL_SORT"
    )
    index_params.add_index(
        field_name="question_vector",
        index_type="IVF_FLAT",
        metric_type="COSINE",
        params={"nlist": 128}
    )
    # 3.5. Create a collection with the index loaded simultaneously
    client.create_collection(
        collection_name="notice",
        schema=schema,
        index_params=index_params
    )
    time.sleep(5)
    res = client.get_load_state(
        collection_name="notice"
    )
    print(res)


def view_collection():
    # 5. View Collections
    client = get_client()
    res = client.describe_collection(
        collection_name="notice"
    )
    print(res)
    res = client.list_collections()
    print(res)
    res_state = client.get_load_state(
        collection_name="customized_setup_1"
    )
    print(res_state)
    client.release_collection("customized_setup_1")
    res_state2 = client.get_load_state(
        collection_name="customized_setup_1"
    )
    print(res_state2)
    res_alias = client.list_aliases(
        collection_name="customized_setup_1"
    )
    print(res_alias)


def add_vector_index(collection_name, field_name="question_vector", metric_type="COSINE", index_type='IVF_FLAT',
                     index_name='vector_index'):
    # 1.prepare index_params
    index_params = MilvusClient.prepare_index_params()

    # 2. Add an index on the vector field.
    index_params.add_index(
        field_name=field_name,
        metric_type=metric_type,
        index_type=index_type,
        index_name=index_name,
        params={"nlist": 128}
    )
    # 3. Create an index file
    get_client().create_index(
        collection_name=collection_name,
        index_params=index_params
    )


def create_collection(collection_name):
    client = get_client()
    if client.has_collection(collection_name):
        drop_collection(collection_name)
    schema = MilvusClient.create_schema(
        auto_id=False,
        enable_dynamic_field=True
    )
    schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
    schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=2560)
    schema.add_field(field_name="router", datatype=DataType.VARCHAR, max_length=2048)
    schema.add_field(field_name="answer", datatype=DataType.VARCHAR, max_length=2048)

    client.create_collection(
        collection_name=collection_name,
        schema=schema,
        auto_id=False
    )
    add_vector_index(collection_name)


def import_data(datas: list):
    count = 0
    client = get_client()
    start_time = time.time()
    # 插入数据
    end_time = time.time()
    try:
        client.insert(collection_name="notice", data=datas)
        count = len(datas)
    except Exception as e:
        log.info('insert has an error={},data={},count={}', e, datas, count)
    log.info(f'import_data success, count={count} ,emb_cost_time={time.time() - start_time}')
    return count


def test_insert_datas():
    embedding_fn = demo.get_embedding
    """
        schema.add_field(field_name="project_name", datatype=DataType.VARCHAR, max_length=128)
        schema.add_field(field_name="project_no", datatype=DataType.VARCHAR, max_length=64)
        schema.add_field(field_name="question_vector", datatype=DataType.FLOAT_VECTOR, dim=768)
        schema.add_field(field_name="answer", datatype=DataType.VARCHAR, max_length=65535)
    """
    to_embedding_text = "联系方式"
    data = {"project_name": "梨树县 2025 年高标准农田建设储备项目勘察设计（第二批）1 标段", "project_no": "NO110",
            "question_vector": embedding_fn([to_embedding_text], convert_to_numpy=False)[0], "answer": """
               招标人：梨树县高标准农田建设项目领导小组办公室
               地 址：梨树县梨树镇树文街 25 号
               联系人：顾明
               联系电话：0434-5288177
               招标代理机构：吉林省嘉丰项目管理有限公司
               地 址：四平市铁西区
               联系人：郭磊
               联系方式：0434-3281189
               监督管理部门：四平市农业农村局"""
            }
    import_data([data])


# def test_embedding():
#     emb_model = WrappedSparkModel.get_spark_emb_model()
#     r = emb_model.embed_query('我是一个兵')
#     print(len(r))

def query_test():
    client = get_client()
    # client.load_collection(collection_name="quick_setup")
    print(client.get_load_state(collection_name="quick_setup"))
    res = get_client().query(
        collection_name="quick_setup",
        # ids=[0, 1, 2],
        filter="subject == 'QA'",
        output_fields=["id", "router", 'answer'],

    )
    print(len(res))
    print(len(res[0]['router']))
    # print(res)


def has_entity(collection_name, project_no):
    res = get_client().query(
        collection_name=collection_name,
        # ids=[0, 1, 2],
        filter=f"project_no == '{project_no}'",
        output_fields=["count(*)"],
    )
    return res[0]['count(*)'] > 0


def count(collection_name):
    res = get_client().query(
        collection_name=collection_name,
        # highlight-start
        output_fields=["count(*)"]
        # highlight-end
    )
    print(res)


def drop_collection(collection_name):
    client = get_client()
    r = client.drop_collection(collection_name)
    print(r)


def del_entities(collection_name, project_no):
    client = get_client()
    r = client.delete(collection_name,
                      filter=f"project_no == '{project_no}'")
    print(r)


def similarity_search(query, topK=2):
    if query:
        emb_model_fn = demo.get_embedding
        query_vector = emb_model_fn([query])[0]
        client = get_client()
        r = client.search(collection_name='notice', data=[query_vector], limit=topK,
                          # search_params={"metric_type": "COSINE", "params": {}},
                          output_fields=["notice_id", "question", "project_no", "project_name", 'answer'])
        context_id_list = []
        context_question_answer_list = []
        if r[0]:
            context_id_list = [result["entity"]["notice_id"] for result in r[0]]
            context_question_answer_list = [
                {"question": result["entity"]["question"], "answer": result["entity"]["answer"]} for result in r[0]]
        log.info(
            "similarity_search topK={},context_id_list={},context_question_answer_list={}",
            topK, context_id_list, context_question_answer_list)
        return context_question_answer_list


if __name__ == '__main__':
    # t1()
    # create_connection()
    custom_schema()
    # view_collection()
    # drop_collection("notice")
    # create_collection()
    # add_vector_index('quick_setup')
    # test_embedding()
    # test_insert_datas()
    # query_test()
    # del_entities('notice', 'SPS20240923NTGC01001')
    # count('quick_setup')

    # similarity_search('联系方式')
    # has_entity("notice", "SPS20240923NTGC01001")
    pass
